from flask import Flask, request, jsonify, render_template, send_from_directory
import os
import subprocess
import tempfile
import json
import logging
import uuid
import shutil
from datetime import datetime
import re
from pathlib import Path
import requests

app = Flask(__name__, static_folder='static', template_folder='templates')

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 配置
UPLOAD_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'static', 'uploads')
RESULTS_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'static', 'results')

# 确保目录存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(RESULTS_FOLDER, exist_ok=True)

# 默认模型路径 - 根据您的环境调整
DEFAULT_MODEL_PATH = "F:\\workspace\\python\\new\\patch2\\model\\MVTecAD_Results\\IM384_WR50_L2-3_P01_D1024-1024_PS-3_AN-1_S0\\models\\mvtec_blade0"

def generate_depth_image(rgb_image_path):
    """调用深度图生成API"""
    url = "http://localhost:5005/generate_depth"
    rgb_image_path = Path(rgb_image_path).as_posix()
    data = {
        "image_path": rgb_image_path
    }
    
    try:
        response = requests.post(url, json=data)
        response.raise_for_status()  # 如果请求失败会抛出异常
        
        # 构建深度图路径
        rgb_filename = os.path.basename(rgb_image_path)
        depth_filename = os.path.splitext(rgb_filename)[0] + '.png'
        depth_path = os.path.join('F:/fan/uploadPath/upload/depth', depth_filename)
        
        # 检查深度图是否生成成功
        if os.path.exists(depth_path):
             # 同步复制到 static 目录
            static_depth_dir = os.path.join(app.static_folder, 'depth')
            os.makedirs(static_depth_dir, exist_ok=True)
            static_depth_path = os.path.join(static_depth_dir, depth_filename)
            shutil.copy(depth_path, static_depth_path)
            return depth_path
        else:
            logger.error(f"深度图生成失败，路径不存在: {depth_path}")
            raise Exception("深度图生成失败")
            
    except Exception as e:
        logger.error(f"调用深度图生成API失败: {str(e)}")
        raise Exception(f"调用深度图生成API失败: {str(e)}")

@app.route('/')
def index():
    """渲染主页"""
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    """处理图像预测请求"""
    try:
        # 检查是否有RGB图像
        if 'rgb_image' not in request.files:
            return jsonify({'error': '未上传RGB图像'}), 400
        
        rgb_file = request.files['rgb_image']
        if rgb_file.filename == '':
            return jsonify({'error': '未选择RGB图像'}), 400
        
        # 创建唯一的会话ID
        session_id = datetime.now().strftime("%Y%m%d_%H%M%S") + "_" + str(uuid.uuid4())[:8]
        
        # 创建会话目录
        session_upload_dir = os.path.join(UPLOAD_FOLDER, session_id)
        session_results_dir = os.path.join(RESULTS_FOLDER, session_id)
        
        os.makedirs(session_upload_dir, exist_ok=True)
        os.makedirs(session_results_dir, exist_ok=True)
        
        # 保存RGB图像
        rgb_filename = os.path.basename(rgb_file.filename)
        rgb_path = os.path.join(session_upload_dir, rgb_filename)
        rgb_file.save(rgb_path)
        
        # 处理深度图
        depth_path = None
        depth_filename = None
        depth_image_path = None  # 用于存储返回给前端的路径
        if 'depth_image' in request.files and request.files['depth_image'].filename != '':
            depth_file = request.files['depth_image']
            depth_filename = os.path.basename(depth_file.filename)
            depth_path = os.path.join(session_upload_dir, depth_filename)
            depth_file.save(depth_path)
            depth_image_path = os.path.join('uploads', session_id, depth_filename)
        else:
            # 如果没有上传深度图，调用API生成
            try:
                # 先检查是否存在预先生成的深度图
                depth_dir = 'F:/fan/uploadPath/upload/depth'
                depth_filename = os.path.splitext(rgb_filename)[0] + '.png'
                preexisting_depth_path = os.path.join(depth_dir, depth_filename)
                
                if os.path.exists(preexisting_depth_path):
                    depth_path = preexisting_depth_path
                    logger.info(f"使用已有深度图: {depth_path}")
                    # 复制到会话目录
                    session_depth_path = os.path.join(session_upload_dir, depth_filename)
                    shutil.copy(depth_path, session_depth_path)
                    depth_path = session_depth_path
                    depth_image_path = os.path.join('uploads', session_id, depth_filename)
                else:
                    # 调用API生成深度图
                    depth_path = generate_depth_image(rgb_path)
                    depth_filename = os.path.basename(depth_path)
                    # 复制到会话目录
                    session_depth_path = os.path.join(session_upload_dir, depth_filename)
                    shutil.copy(depth_path, session_depth_path)
                    depth_path = session_depth_path
                    depth_image_path = os.path.join('uploads', session_id, depth_filename)
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        # 获取其他参数
        model_path = request.form.get('model_path', DEFAULT_MODEL_PATH)
        resize = request.form.get('resize', '512')
        imagesize = request.form.get('imagesize', '384')
        
        # 构建命令
        cmd = [
            "python", "bin/load_and_evaluate_patchcore_s.py",
            "--gpu", "0",
            "--seed", "0",
            session_results_dir,
            "patch_core_loader",
            "-p", model_path,
            "single_image",
            "--resize", resize,
            "--imagesize", imagesize,
            rgb_path,
            depth_path
        ]
        
        # 执行命令
        logger.info(f"执行命令: {' '.join(cmd)}")
        process = subprocess.run(cmd, capture_output=True, text=True, encoding='utf-8')
        
        # 尝试从输出中提取异常分数
        anomaly_score = 0
        score_match = re.search(r"Anomaly score: ([0-9.]+)", process.stdout)
        if score_match:
            anomaly_score = float(score_match.group(1))
        
        if process.returncode != 0:
            logger.error(f"命令执行失败: {process.stderr}")
            return jsonify({
                'warning': '模型执行部分失败，但成功提取了异常分数',
                'anomaly_score': anomaly_score,
                'rgb_image': os.path.join('uploads', session_id, rgb_filename),
                'depth_image': depth_image_path,
                'segmentation_images': [],
                'timestamp': datetime.now().isoformat()
            })
        
        # 读取结果
        try:
            result_file = os.path.join(session_results_dir, "single_image_result.json")
            if os.path.exists(result_file):
                with open(result_file, 'r') as f:
                    result = json.load(f)
                    anomaly_score = result.get('anomaly_score', anomaly_score)
            else:
                logger.warning(f"结果文件不存在: {result_file}，使用从输出提取的分数")
            
            # 获取分割图路径（如果存在）
            segmentation_images = []
            re_dir = os.path.join(session_results_dir, "re")
            if os.path.exists(re_dir):
                for root, _, files in os.walk(re_dir):
                    for file in files:
                        if file.endswith(('.png', '.jpg')):
                            rel_path = os.path.relpath(os.path.join(root, file), start=app.static_folder)
                            segmentation_images.append(rel_path)
            
            # 构建响应
            response = {
                'session_id': session_id,
                'anomaly_score': anomaly_score,
                'rgb_image': os.path.join('uploads', session_id, rgb_filename),
                'depth_image': depth_image_path,
                'segmentation_images': segmentation_images,
                'timestamp': datetime.now().isoformat()
            }
            
            return jsonify(response)
            
        except Exception as e:
            logger.exception("读取结果时出错")
            return jsonify({
                'warning': '读取结果失败，但成功提取了异常分数',
                'anomaly_score': anomaly_score,
                'rgb_image': os.path.join('uploads', session_id, rgb_filename),
                'depth_image': depth_image_path,
                'segmentation_images': [],
                'timestamp': datetime.now().isoformat()
            })
        
    except Exception as e:
        logger.exception("意外错误")
        return jsonify({'error': '意外错误', 'details': str(e)}), 500

@app.route('/results/<path:filename>')
def results(filename):
    """提供结果文件访问"""
    return send_from_directory(RESULTS_FOLDER, filename)

@app.route('/health', methods=['GET'])
def health_check():
    """健康检查端点"""
    return jsonify({'status': 'ok'})

@app.route('/detect', methods=['POST'])
def detect():
    """处理单张RGB图像的深度图生成和异常检测"""
    try:
        # 检查是否有RGB图像
        if 'rgb_image' not in request.files:
            return jsonify({'error': '未上传RGB图像'}), 400
        
        rgb_file = request.files['rgb_image']
        if rgb_file.filename == '':
            return jsonify({'error': '未选择RGB图像'}), 400
        
        # 创建临时目录保存图像
        temp_dir = os.path.join(app.static_folder, 'temp')
        os.makedirs(temp_dir, exist_ok=True)
        
        # 保存RGB图像
        rgb_filename = os.path.basename(rgb_file.filename)
        rgb_path = os.path.join(temp_dir, rgb_filename)
        rgb_file.save(rgb_path)
        
        # 优先检查是否已存在深度图
        depth_dir = 'F:/fan/uploadPath/upload/depth'
        depth_filename = os.path.splitext(rgb_filename)[0] + '.png'
        preexisting_depth_path = os.path.join(depth_dir, depth_filename)

        if os.path.exists(preexisting_depth_path):
            depth_path = preexisting_depth_path
            logger.info(f"使用已有深度图: {depth_path}")
            # 复制到静态目录以便返回给前端
            static_depth_dir = os.path.join(app.static_folder, 'depth')
            os.makedirs(static_depth_dir, exist_ok=True)
            static_depth_path = os.path.join(static_depth_dir, depth_filename)
            shutil.copy(depth_path, static_depth_path)
            # 用于返回给前端的相对路径
            depth_image_path = os.path.join('depth', depth_filename)
        else:
            # 调用API生成深度图
            try:
                depth_path = generate_depth_image(rgb_path)
                depth_filename = os.path.basename(depth_path)
                # 用于返回给前端的相对路径
                depth_image_path = os.path.join('depth', os.path.basename(depth_path))
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        # 获取其他参数
        model_path = request.form.get('model_path', DEFAULT_MODEL_PATH)
        resize = request.form.get('resize', '512')
        imagesize = request.form.get('imagesize', '384')
        
        # 创建结果目录
        results_dir = os.path.join(app.static_folder, 'results', 'temp')
        os.makedirs(results_dir, exist_ok=True)
        
        # 构建命令
        cmd = [
            "python", "bin/load_and_evaluate_patchcore_s.py",
            "--gpu", "0",
            "--seed", "0",
            results_dir,
            "patch_core_loader",
            "-p", model_path,
            "single_image",
            "--resize", resize,
            "--imagesize", imagesize,
            rgb_path,
            depth_path
        ]
        
        # 执行命令
        logger.info(f"执行命令: {' '.join(cmd)}")
        process = subprocess.run(cmd, capture_output=True, text=True, encoding='utf-8')
        
        # 尝试从输出中提取异常分数
        anomaly_score = 0
        score_match = re.search(r"Anomaly score: ([0-9.]+)", process.stdout)
        if score_match:
            anomaly_score = float(score_match.group(1))
        
        if process.returncode != 0:
            logger.error(f"命令执行失败: {process.stderr}")
            return jsonify({
                'warning': '模型执行部分失败，但成功提取了异常分数',
                'anomaly_score': anomaly_score
            })
        
        # 读取结果文件
        try:
            result_file = os.path.join(results_dir, "single_image_result.json")
            if os.path.exists(result_file):
                with open(result_file, 'r') as f:
                    result = json.load(f)
                    anomaly_score = result.get('anomaly_score', anomaly_score)
        except Exception as e:
            logger.warning(f"读取结果文件时出错: {str(e)}")
        
        # 清理临时文件
        # try:
        #     os.remove(rgb_path)
        #     os.remove(depth_path)
        #     shutil.rmtree(results_dir)
        # except Exception as e:
        #     logger.warning(f"清理临时文件时出错: {str(e)}")
        
        return jsonify({
            'anomaly_score': anomaly_score,
            'status': 'success'
        })
        
    except Exception as e:
        logger.exception("处理过程中发生错误")
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)